import pandas as pd
import tushare as ts
import numpy as np
import time
import datetime


#求均线
def ma(day):
    df = pd.read_hdf('stock_data.h5', key='data')
    data = df.swaplevel('trade_date', 'ts_code')

    data.drop('ts_code', axis=1, inplace=True)  # 数据筛选
    data2 = data.reset_index('trade_date')
    data2 = data.sort_index(level=0)

    stock_list = data2.index.values
    stock_list = list(set(stock_list))

    df_new = pd.DataFrame()

    dict = {}
    i=0
    for stock in stock_list:

        i=i+1
        print(i)
        print('正在进行：',stock)
        ma_day = round(data2[stock]['close'].rolling(day).mean(), 2)
        ewm_day = round(data2[stock]['close'].ewm(day).mean(), 2)
        #字典
        dict['ts_code'] = stock
        dict['MA_' + str(day)] = ma_day
        dict['EMA_' + str(day)] = ewm_day
        #数据存入dataframe
        df1 = pd.DataFrame(dict)
        df_new = df_new.append(df1)
    print(df_new)

start = time.clock()
ma(5)
end = time.clock()
print('程序耗时：',end - start)
